from util_for_argparse import get_cuda_port_args

args = get_cuda_port_args()

from flask import Flask, request, jsonify, send_file
from io import BytesIO
from PIL import Image

from util_parse import get_parse_predictor, get_parse_res

app = Flask(__name__)

# Build the parse predictor at startup for efficiency
parse_predictor = get_parse_predictor()

@app.route('/parse', methods=['POST'])
def parse_image():
    if 'image' not in request.files:
        return jsonify({'error': 'No image uploaded'}), 400
    file = request.files['image']
    try:
        img = Image.open(file.stream).convert("RGB")
    except Exception as e:
        return jsonify({'error': f'Error loading image: {str(e)}'}), 400
    try:
        result = get_parse_res(parse_predictor=parse_predictor, img_pil=img)
        # Result is PIL image; serve directly
        buf = BytesIO()
        result.save(buf, format='PNG')
        buf.seek(0)
        return send_file(
            buf, mimetype='image/png', as_attachment=False, download_name='parse.png'
        )
    except Exception as e:
        return jsonify({'error': f'Parse error: {str(e)}'}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=args.port, debug=False)
